﻿using Ndns.ServerHost.Abstractions;
using System;
using Ndns.ServerHost.Abstractions.Services;
using System.Threading.Tasks;
using System.Collections.Generic;
using Ndns.Packet.Abstractions;
using System.IO;
using System.Net;
using System.Net.Sockets;
using Ndns.Packet.Abstractions.ResourceRecords;
using System.Linq;

namespace Ndns.HostsFile
{
    public class HostsFileMiddleware : IMiddleware
    {
        private readonly string _fileName;
        private readonly Dictionary<DnsHostName, List<IPAddress>> _v4hosts = new Dictionary<DnsHostName, List<IPAddress>>();
        private readonly Dictionary<DnsHostName, List<IPAddress>> _v6hosts = new Dictionary<DnsHostName, List<IPAddress>>();

        private RequestDelegate _next;
        private ILogger _logger;

        private readonly HostsFileMiddlewareOptions _options;

        public HostsFileMiddleware(string fileName, HostsFileMiddlewareOptions options)
        {
            _fileName = fileName;
            _options = options;
        }

        public void RegisterServices(IServiceCollection services, RequestDelegate next)
        {
            _next = next;
            _logger = services.Get<ILoggerProvider>().GetLogger<HostsFileMiddleware>();
        }

        public void Start()
        {
            InitDictionaryFromFile();
        }

        private void InitDictionaryFromFile()
        {
            try
            {
                using (var fs = new FileStream(_fileName, FileMode.Open, FileAccess.Read))
                using (var sr = new StreamReader(fs))
                {
                    while (!sr.EndOfStream)
                    {
                        string line = sr.ReadLine();
                        line = line.Replace('\t', ' ').Trim();
                        if (line.StartsWith("#", StringComparison.OrdinalIgnoreCase))
                        {
                            continue;
                        }

                        while (line.Contains("  "))
                        {
                            line = line.Replace("  ", " ");
                        }

                        var stringArray = line.Split(' ');
                        if (stringArray.Length != 2)
                        {
                            continue;
                        }

                        if (!IPAddress.TryParse(stringArray[0], out var ipaddr))
                        {
                            continue;
                        }

                        var hostName = new DnsHostName(stringArray[1]);

                        Dictionary<DnsHostName, List<IPAddress>> targetDict = null;
                        if (ipaddr.AddressFamily == AddressFamily.InterNetwork)
                        {
                            targetDict = _v4hosts;
                        }
                        else if (ipaddr.AddressFamily == AddressFamily.InterNetworkV6)
                        {
                            targetDict = _v6hosts;
                        }

                        if (targetDict == null)
                        {
                            continue;
                        }

                        if (targetDict.ContainsKey(hostName))
                        {
                            targetDict[hostName].Add(ipaddr);
                        }
                        else
                        {
                            targetDict.Add(hostName, new List<IPAddress>
                        {
                            ipaddr
                        });
                        }
                    }
                }

                _logger.LogInfo(string.Empty, (_v4hosts.Count + _v6hosts.Count) + $" record(s) read from hosts file. ({_fileName})");
            }
            catch (Exception ex)
            {
                _v4hosts.Clear();
                _v6hosts.Clear();
                
                _logger.LogError(string.Empty, $"Failed to read records from hosts file. ({_fileName})", ex);
            }
            
        }

        public Task InvokeAsync(DnsContext context)
        {
            var request = context.Request;
            if (request.Queries.Count != 1)
            {
                _logger.LogInfo(context.TraceIdentifier, "Can't handle this request because it contains more than one query");

                return _next(context);
            }

            var query = request.Queries[0];
            if (query.Type == DnsRecordType.A)
            {
                return HandleV4Request(context, query);
            }
            else if (query.Type == DnsRecordType.AAAA)
            {
                return HandleV6Request(context, query);
            }

            _logger.LogTrace(context.TraceIdentifier, "Can't handle this request because record type is unsupported.");

            return _next(context);
        }

        private async Task HandleV4Request(DnsContext context, DnsQuery query)
        {
            var response = context.Response;

            if (_v4hosts.ContainsKey(query.HostName))
            {
                var list = _v4hosts[query.HostName];

                _logger.LogInfo(context.TraceIdentifier, $"{list.Count} record(s) found in the hosts file.");

                foreach (var item in list)
                {
                    response.AnswerRRs.Add(new DnsResourceRecordA(query.HostName, query.Class, _options.DefaultTTL, item));
                }
                response.ReplyCode = DnsReplyCode.NoError;
                response.ServerAuthoritative = false;
                response.RecursionAvailable = _options.IsRecursionAvailable;
                response.AnswerAuthenticated = false;
                return;
            }
            
            _logger.LogTrace(context.TraceIdentifier, $"No record found in the hosts file.");

            await _next(context);
            
            if (response.ReplyCode != DnsReplyCode.NoError || response.IgnoreRequest || response.AnswerRRs.Count == 0)
            {
                return;
            }
            
            if (_options.EnableCNAMEResolution)
            {
                _logger.LogTrace(context.TraceIdentifier, $"Trying to resolve CNAME records.");

                DnsHostName currentHostName = query.HostName;
                var records = new List<DnsResourceRecord>();

                do
                {
                    DnsResourceRecordCNAME currentRecord = response.AnswerRRs.Where(r => r.HostName == currentHostName).OfType<DnsResourceRecordCNAME>().FirstOrDefault();
                    if (currentRecord == null)
                    {
                        return;
                    }

                    records.Add(currentRecord);
                    currentHostName = currentRecord.CanonicalName;

                    if (_v4hosts.ContainsKey(currentHostName))
                    {
                        var list = _v4hosts[currentHostName];

                        _logger.LogInfo(context.TraceIdentifier, $"CNAME records resolve to {currentHostName} {list.Count} record(s) found in the hosts file.");

                        foreach (var item in list)
                        {
                            records.Add(new DnsResourceRecordA(query.HostName, query.Class, _options.DefaultTTL, item));
                        }

                        break;
                    }
                } while (true);

                response.AnswerRRs.Clear();
                response.AuthoritativeNSs.Clear();
                response.AdditionalRRs.Clear();

                foreach (var item in records)
                {
                    response.AnswerRRs.Add(item);
                }

                response.ReplyCode = DnsReplyCode.NoError;
                response.ServerAuthoritative = false;
                response.RecursionAvailable = _options.IsRecursionAvailable;
                response.AnswerAuthenticated = false;
            }
        }

        private async Task HandleV6Request(DnsContext context, DnsQuery query)
        {
            var response = context.Response;

            if (_v6hosts.ContainsKey(query.HostName))
            {
                var list = _v6hosts[query.HostName];

                _logger.LogInfo(context.TraceIdentifier, $"{list.Count} record(s) found in the hosts file.");

                foreach (var item in list)
                {
                    response.AnswerRRs.Add(new DnsResourceRecordAAAA(query.HostName, query.Class, _options.DefaultTTL, item));
                }
                response.ReplyCode = DnsReplyCode.NoError;
                response.ServerAuthoritative = false;
                response.RecursionAvailable = _options.IsRecursionAvailable;
                response.AnswerAuthenticated = false;
                return;
            }

            _logger.LogTrace(context.TraceIdentifier, $"No record found in the hosts file.");

            await _next(context);

            if (response.ReplyCode != DnsReplyCode.NoError || response.IgnoreRequest || response.AnswerRRs.Count == 0)
            {
                return;
            }

            if (_options.EnableCNAMEResolution)
            {
                _logger.LogTrace(context.TraceIdentifier, $"Trying to resolve CNAME records.");

                DnsHostName currentHostName = query.HostName;
                var records = new List<DnsResourceRecord>();

                do
                {
                    DnsResourceRecordCNAME currentRecord = response.AnswerRRs.Where(r => r.HostName == currentHostName).OfType<DnsResourceRecordCNAME>().FirstOrDefault();
                    if (currentRecord == null)
                    {
                        return;
                    }

                    records.Add(currentRecord);
                    currentHostName = currentRecord.CanonicalName;

                    if (_v6hosts.ContainsKey(currentHostName))
                    {
                        var list = _v6hosts[currentHostName];

                        _logger.LogInfo(context.TraceIdentifier, $"CNAME records resolve to {currentHostName} {list.Count} record(s) found in the hosts file.");

                        foreach (var item in list)
                        {
                            records.Add(new DnsResourceRecordAAAA(query.HostName, query.Class, _options.DefaultTTL, item));
                        }

                        break;
                    }
                } while (true);

                response.AnswerRRs.Clear();
                response.AuthoritativeNSs.Clear();
                response.AdditionalRRs.Clear();

                foreach (var item in records)
                {
                    response.AnswerRRs.Add(item);
                }

                response.ReplyCode = DnsReplyCode.NoError;
                response.ServerAuthoritative = false;
                response.RecursionAvailable = _options.IsRecursionAvailable;
                response.AnswerAuthenticated = false;
            }
        }

        #region IDisposable Support
        private bool disposedValue = false; // 要检测冗余调用

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // 释放托管状态(托管对象)。
                }

                // 释放未托管的资源(未托管的对象)并在以下内容中替代终结器。
                // 将大型字段设置为 null。

                disposedValue = true;
            }
        }

        // 仅当以上 Dispose(bool disposing) 拥有用于释放未托管资源的代码时才替代终结器。
        // ~HostsFileMiddleware() {
        //   // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
        //   Dispose(false);
        // }

        // 添加此代码以正确实现可处置模式。
        public void Dispose()
        {
            // 请勿更改此代码。将清理代码放入以上 Dispose(bool disposing) 中。
            Dispose(true);
            // 如果在以上内容中替代了终结器，则取消注释以下行。
            // GC.SuppressFinalize(this);
        }
        #endregion
    }
}
